/* Copyright (c) 2024 Huawei Technologies Co., Ltd.
* This file is a part of the CANN Open Software.
* Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
* Please refer to the License for details. You may not use this file except in compliance with the License.
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
* INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
* See LICENSE in the root of the software repository for the full text of the License.
* ===================================================================================================================*/
#include <gtest/gtest.h>
#include "ascend_ir/expression.h"
#include <iostream>
#include <cmath>

#include <symengine/integer.h>
#include <symengine/assumptions.h>
#include <symengine/functions.h>
#include <symengine/simplify.h>
#include <symengine/integer.h>
#include <symengine/real_double.h>

class UtestExpression : public testing::Test {
 protected:
  void SetUp() {}

  void TearDown() {}
};
using namespace ge::expression;

// // // symengine error1
// TEST_F(UtestExpression, Symengine_error1) {
//   auto e1 = SymEngine::integer(3);
//   auto e4 = SymEngine::symbol("x");
//   auto e_pow = SymEngine::pow(e4, e1);
//   auto e5 = SymEngine::div(SymEngine::one, e_pow);
//   auto e6 = e5->get_args();
//   for (auto &arg : e6) {
//     std::cout << arg->__str__() << std::endl;
//   }
// }

// // // symengine error2
// TEST_F(UtestExpression, Symengine_error2) {
//   auto neg_x = SymEngine::mul(SymEngine::symbol("x"), SymEngine::minus_one);
// }

// TEST_F(UtestExpression, Symengine) {
//   // for (auto &arg : SymEngine::sub(SymEngine::symbol("x"), SymEngine::symbol("y"))->get_args()) {
//   //   std::cout << arg->__str__() << std::endl;
//   //   for (auto &arg2 : arg->get_args()) {
//   //     std::cout << arg2->__str__() << std::endl;
//   //   }
//   // }

//   auto x = SymEngine::add(SymEngine::symbol("x"), SymEngine::integer(1));
//   bool has = false;
//   for (auto &arg : x->get_args()) {
//     has = true;
//     std::cout << "arg:" << arg->__str__() << std::endl;
//   }
//   if (!has) {
//     std::cout << x->__str__() << std::endl;
//   }


// }

TEST_F(UtestExpression, GetExprType) {
  auto expr1 = Expression::CreateExpression(2);
  EXPECT_NE(expr1, nullptr);
  EXPECT_EQ(expr1->GetExprType(), ExprType::EXPR_CONSTANT);

  auto expr_uint32 = Expression::CreateExpression(2u);
  EXPECT_NE(expr_uint32, nullptr);
  EXPECT_EQ(expr_uint32->GetExprType(), ExprType::EXPR_CONSTANT);

  auto expr_int64 = Expression::CreateExpression(2l);
  EXPECT_NE(expr_int64, nullptr);
  EXPECT_EQ(expr_int64->GetExprType(), ExprType::EXPR_CONSTANT);

  auto expr_uint64 = Expression::CreateExpression(2lu);
  EXPECT_NE(expr_uint64, nullptr);
  EXPECT_EQ(expr_uint64->GetExprType(), ExprType::EXPR_CONSTANT);

  auto expr2 = Expression::CreateExpression("a");
  EXPECT_NE(expr2, nullptr);
  EXPECT_EQ(expr2->GetExprType(), ExprType::EXPR_VARIABLE);

  auto expr3 = Add(expr1, expr2);
  EXPECT_NE(expr3, nullptr);
  EXPECT_EQ(expr3->GetExprType(), ExprType::EXPR_OPERATION);
  EXPECT_EQ(expr3 == expr2, false);
  EXPECT_EQ(*expr3 == *expr2, false);

  expr3 = Sub(expr1, expr2);
  EXPECT_NE(expr3, nullptr);
  EXPECT_EQ(expr3->GetExprType(), ExprType::EXPR_OPERATION);

  expr3 = Mul(expr1, expr2);
  EXPECT_NE(expr3, nullptr);
  EXPECT_EQ(expr3->GetExprType(), ExprType::EXPR_OPERATION);

  expr3 = Div(expr1, expr2);
  EXPECT_NE(expr3, nullptr);
  EXPECT_EQ(expr3->GetExprType(), ExprType::EXPR_OPERATION);

  expr3 = Max(expr1, expr2);
  EXPECT_NE(expr3, nullptr);
  EXPECT_EQ(expr3->GetExprType(), ExprType::EXPR_OPERATION);

  expr3 = Min(expr1, expr2);
  EXPECT_NE(expr3, nullptr);
  EXPECT_EQ(expr3->GetExprType(), ExprType::EXPR_OPERATION);

  expr3 = Pow(expr1, expr2);
  EXPECT_NE(expr3, nullptr);
  EXPECT_EQ(expr3->GetExprType(), ExprType::EXPR_OPERATION);

  expr3 = Log(expr1, expr2);
  EXPECT_NE(expr3, nullptr);
  EXPECT_EQ(expr3->GetExprType(), ExprType::EXPR_OPERATION);

  expr3 = Log(expr2);
  EXPECT_NE(expr3, nullptr);
  EXPECT_EQ(expr3->GetExprType(), ExprType::EXPR_OPERATION);

  expr3 = Ceiling(expr2);
  EXPECT_NE(expr3, nullptr);
  EXPECT_EQ(expr3->GetExprType(), ExprType::EXPR_OPERATION);

  expr3 = Rational(1, 2);
  EXPECT_NE(expr3, nullptr);
  EXPECT_EQ(expr3->GetExprType(), ExprType::EXPR_CONSTANT);
}

TEST_F(UtestExpression, Str) {
  auto var_b = Expression::CreateExpression("b");
  auto var_c = Expression::CreateExpression("c");
  auto b_div_c = Div(var_b, var_c);
  EXPECT_EQ(b_div_c->Str(), "(b / (c))");

  EXPECT_EQ(Rational(2, 3)->Str(), "(double)(2)/(double)(3)");
  EXPECT_EQ(Pow(var_b, Rational(1, 2))->Str(), "sqrt(b)");
  EXPECT_EQ(Pow(var_b, Expression::CreateExpression(3))->Str(), "(b * b * b)");
  EXPECT_EQ(Pow(var_b, var_c)->Str(), "pow(b, c)");
  EXPECT_EQ(Ceiling(var_b)->Str(), "ceiling(b)");
  EXPECT_EQ(Sub(Expression::CreateExpression(2), Expression::CreateExpression(3) * var_b)->Str(), "(2 - (3 * b))");
  EXPECT_EQ(Max(Max(var_b, var_c), Expression::CreateExpression("d"))->Str(),
    "Max(Max(b, c), d)");
  EXPECT_EQ(Min(Min(var_b, var_c), Expression::CreateExpression("d"))->Str(),
    "Min(Min(b, c), d)");
}

TEST_F(UtestExpression, Operator) {
  auto var_b = Expression::CreateExpression("b");
  auto var_c = Expression::CreateExpression("c");
  EXPECT_EQ((var_b + var_c)->GetExprType(), ExprType::EXPR_OPERATION);
  EXPECT_EQ((var_b - var_c)->GetExprType(), ExprType::EXPR_OPERATION);
  EXPECT_EQ((var_b * var_c)->GetExprType(), ExprType::EXPR_OPERATION);
  EXPECT_EQ((var_b / var_c)->GetExprType(), ExprType::EXPR_OPERATION);
}

TEST_F(UtestExpression, Replace) {
  auto const1 = Expression::CreateExpression(1, "const1");
  EXPECT_EQ(const1->GetName(), "const1");
  auto var_b = Expression::CreateExpression("b");
  EXPECT_EQ(var_b->GetName(), "b");
  auto var_c = Expression::CreateExpression("c");
  auto var_d = Expression::CreateExpression("d");
  auto b_div_c = Div(var_b, var_c);
  std::map<ExpressionPtr, ExpressionPtr> replace_vars;
  replace_vars[var_b] = var_d;
  auto replace_expr = b_div_c->Replace(replace_vars);
  EXPECT_TRUE(replace_expr->Eq(Div(var_d, var_c)));
  EXPECT_FALSE(replace_expr->Eq(nullptr));
}

TEST_F(UtestExpression, Simplify) {
  auto var_b = Expression::CreateExpression("b");
  auto const_1 = Expression::CreateExpression(1);
  auto const_2 = Expression::CreateExpression(2);
  auto const_3 = Expression::CreateExpression(3);
  EXPECT_TRUE(*(Add(Add(var_b, const_1), const_2)->Simplify()) == *Add(var_b, const_3));
}

TEST_F(UtestExpression, GetPrimaryArgs) {
  auto const_neg_2 = Expression::CreateExpression(-2);
  auto const_neg_3 = Expression::CreateExpression(-3);
  auto var_b = Expression::CreateExpression("b");
  auto var_c = Expression::CreateExpression("c");
  auto var_d = Expression::CreateExpression("d");
  auto var_e = Expression::CreateExpression("e");
  std::vector<ExpressionPtr> args_exp = {var_b, var_c, var_d, var_e};
  auto mul_expr = Min(Max(Add(Pow(Mul(const_neg_2, var_b), Mul(var_b, const_neg_3)), var_c), var_d), var_e);
  std::cout << *mul_expr << std::endl;
  auto prim_args = mul_expr->GetPrimaryArgs();
  EXPECT_EQ(prim_args.size(), args_exp.size());
  bool has_find = true; 
  for (auto &arg_get : prim_args) {
    bool one_has_find = false;
    for (auto &arg_exp : args_exp) {
      if (*arg_get == *arg_exp) {
        one_has_find = true;
        break;
      }
    }
    if (!one_has_find) {
      has_find = false;
      break;
    }
  }
  EXPECT_EQ(has_find, true);
}

TEST_F(UtestExpression, ContainVar) {
  auto var_b = Expression::CreateExpression("b");
  auto const_1 = Expression::CreateExpression(1);
  EXPECT_TRUE(Add(var_b, const_1)->ContainVar(var_b));
  EXPECT_FALSE(Add(var_b, const_1)->ContainVar(const_1));
}

TEST_F(UtestExpression, GetResult) {
  auto var_b = Expression::CreateExpression("b");
  auto var_c = Expression::CreateExpression("c");
  auto b_add_c = Add(var_b, var_c);
  std::map<ExpressionPtr, ExpressionPtr> replace_vars;
  replace_vars[var_b] = Expression::CreateExpression(1);
  replace_vars[var_c] = Expression::CreateExpression(2);
  double result;
  auto code = b_add_c->GetResult(replace_vars, result);
  EXPECT_EQ(code, ge::GRAPH_SUCCESS);
  EXPECT_EQ(result, static_cast<double>(3));

  replace_vars.clear();
  replace_vars[var_b] = Expression::CreateExpression(1.0);
  replace_vars[var_c] = Expression::CreateExpression(2.0);
  code = b_add_c->GetResult(replace_vars, result);
  EXPECT_EQ(code, ge::GRAPH_SUCCESS);
  EXPECT_EQ(result, static_cast<double>(3));

  replace_vars.clear();
  replace_vars[var_b] = Expression::CreateExpression(1);
  replace_vars[var_c] = Rational(2, 3);
  code = b_add_c->GetResult(replace_vars, result);
  EXPECT_EQ(code, ge::GRAPH_SUCCESS);
  EXPECT_TRUE(std::abs(result - (static_cast<double>(1) + static_cast<double>(2) / static_cast<double>(3))) < 0.0001);
}
